# conding: utf-8
#!/usr/local/python36/bin/python3.6

from threading import Thread
from datetime import datetime

import logging.config, logging
import re, queue, os, time, sys, getopt

#日志
log_dir = "./log"
if os.path.exists(log_dir) == False:
    os.mkdir(log_dir)
logging.config.fileConfig("logging.conf")
trafficlog = logging.getLogger('traffic')

queuelist = queue.Queue()
ssPidJson = {}
ssAddressJson = {}

class DealSSdata(Thread):

    def __init__(self):
        Thread.__init__(self)
        self.regexport = re.compile(':')

    def genData(self):
        data = queuelist.get()
        genlist = []
        for i in data:
            pdata = self.regexport.split(i[4])
            if len(pdata) != 2:
                port = i[5]
            else:
                port = pdata[-1]
            datajson = {"ib": int(i[2]), "ob": int(i[3]), "port": port, "pid": int(i[-1])}
            genlist.append(datajson)
        return genlist

    def pidData(self, genlist):
        for i in genlist:
            key = i["pid"] + "_" + i["port"]
            data = ssPidJson.get(key)
            if data == None:
                i["date"] = datetime.now()
                ssPidJson[key] = i
            else:
                data["ib"] += i["ib"]
                data["ob"] += i["ob"]
                data["date"] = datetime.now()

    def run(self):
        while True:
            if queuelist.empty() == True:
                time.sleep(0.1)
                continue
            data = self.genData()
            self.pidData(data)

class ShowSSdata(Thread):
    def __init__(self, sleeptime, datatype):
        Thread.__init__(self)
        self.sleep_time = sleeptime if sleeptime is not None else 60
        self.deal_time = 3600 #定时清理未修改数据
        self.data_size = {"GB":8589934592,"MB":8388608,"KB":8192}
        self.data_type = datatype if datatype is not None else "KB"
        self.sizedata = 8192

    def printpidData(self):
        if ssPidJson:
            sortdata = sorted(ssPidJson.values(), key=lambda item: item["ob"], reverse=True)
            for d in sortdata:
                _log_data = "进程号: %s,  端口号: %s, 累计入网流量(%s): %s, 累计出网流量(%s): %s" % (
                    d["pid"], d["port"],self.data_type, d["ib"] // self.sizedata, self.data_type, d["ob"] // self.sizedata
                )
                print(_log_data)
                trafficlog.info(_log_data)

    def delpidData(self):
        del_key = []
        for k in ssPidJson.keys():
            if time.time() - time.mktime(ssPidJson[k]["date"].timetuple()) > self.deal_time:
                del_key.append(k)
        trafficlog.info("清理过期数据，编号(PID_PORT): %s" % del_key)
        for k in del_key:
            ssPidJson.pop(k)

    def run(self):
        self.sizedata = self.data_size[self.data_type]
        while True:
            self.delpidData()
            self.printpidData()
            time.sleep(self.sleep_time)


class GetSSdata(Thread):
    def __init__(self):
        Thread.__init__(self)
        self.thread_counts = 5
        self.workerList = list()
        self.regex = re.compile('\s+')

    def Readdata(self):
        data = []
        ssdata = os.popen('ss -nbop | grep ESTAB').readlines()
        for i in ssdata:
            nsi = []
            si = self.regex.split(i)
            pid_group = re.search(r'.*pid=(.*),fd.*', i, re.M | re.I)
            pid_data = 0
            if pid_group != None:
                pid_data = pid_group.group(1)
            nsi += si[0:6]
            nsi.append(pid_data)
            if int(nsi[2]) >= 8192 or int(nsi[3]) >= 8192:
                data.append(nsi)
        if data != []:
            queuelist.put(data)

    def Worker_start(self):
        for i in range(self.thread_counts):
            _DealSSdata = DealSSdata()
            _DealSSdata.start()
            self.workerList.append(_DealSSdata)

    def Print_start(self, sleeptime, sizetype):
        _ShowSSdata = ShowSSdata(sleeptime, sizetype)
        _ShowSSdata.start()
        self.workerList.append(_ShowSSdata)

    def run(self):
        while True:
            self.Readdata()
            # _log_str = "当前队列堆积数量：" + str(queuelist.qsize())
            # print(_log_str)
            time.sleep(1)

def usage():
    print("Usage:[-t(30s-120s)|-s(GB|MB|KB)|-h] [--help] args....")


if __name__ == "__main__":
    sleeptime = None
    sizetype = None
    try:
        shortargs = 't:s:h'
        longargs = ['help']
        opts, args = getopt.getopt(sys.argv[1:], shortargs, longargs)
        for opt, arg in opts:
            if opt in ["-h", "--help"]:
                usage()
                sys.exit(1)
            elif opt == "-t":
                try:
                    if 30 <= int(arg) <= 120:
                        sleeptime = int(arg)
                    else:
                        print("时间区间30s-120s之间")
                        trafficlog.error("时间区间30s-120s之间")
                except:
                    print("-t参数必须是数值")
                    trafficlog.error("-t参数必须是数值")
            elif opt == "-s":
                if arg in ["GB", "MB", "KB"]:
                    sizetype = arg
                else:
                    print("值类型只能是GB|MB|KB")
                    trafficlog.error("值类型只能是GB|MB|KB")
            else:
                print("参数异常")
                trafficlog.error("参数异常")
                usage()
                sys.exit(1)
    except getopt.GetoptError as b:
        print("参数异常")
        trafficlog.error("参数异常")
        usage()
        sys.exit(1)

    a = GetSSdata()
    a.Worker_start()
    a.Print_start(sleeptime, sizetype)
    a.start()
    trafficlog.info("监控启动")